import os

import numpy as np
import torch
import torch.utils.data as tud


class TimeDataset(tud.Dataset):
    def __init__(self, path, input_size=1, windowsize=1, step=1,
                 mode='train', test_size=0, timemode=False, label=False, reverse_label=False):
        """
        可以将csv文件批量转成tensor，包括文件读取、数据标准化、训练集与测试集划分、自回归数据集划分、tensor类型转化
        csv格式：首行为列名，末列为标签，如含时间戳则放在第一列
        :param path: 文件夹路径。该文件夹可包含多个csv文件
        :param input_size: 将多个时间点的输入数据合成为一个输入，该变量表示取的时间点个数，默认为1
        :param windowsize: 将多个时间点的输出数据合成为一个输出，该变量表示取的时间点个数，默认为1
        :param step: 隔多远采样一次数据
        :param mode: train or test
        :param test_size: test 数据集所占比例
        :param timemode: 第一列是否用时间戳。默认为False
        """
        self.time_feature = input_size
        self.windowsize = windowsize
        self.step = step
        self.test_size = test_size
        self.mode = mode
        self.timemode = timemode
        self.inputs, self.labels, self.outputs = self.parse_data(
            path, label, reverse_label=reverse_label)
        self.train_inputs, self.train_labels, self.train_outputs, self.\
            test_inputs, self.test_labels, self.test_outputs = self.split(
                self.inputs, self.labels, self.outputs)
        self.train_inputs = torch.Tensor(self.train_inputs).to(torch.float32)
        self.train_labels = torch.Tensor(self.train_labels).to(torch.float32)
        self.train_outputs = torch.Tensor(self.train_outputs).to(torch.float32)
        self.test_inputs = torch.Tensor(self.test_inputs).to(torch.float32)
        self.test_labels = torch.Tensor(self.test_labels).to(torch.float32)
        self.test_outputs = torch.Tensor(self.test_outputs).to(torch.float32)
        # if not in the test mode
        if test_size < 1:
            self.feature_size = self.train_inputs.shape[-1]
            self.time_feature = self.feature_size * self.train_inputs.shape[1]
        elif test_size == 1:
            self.feature_size = self.test_inputs.shape[-1]
            self.time_feature = self.feature_size * self.test_inputs.shape[1]

    def parse_data(self, path, labeled, reverse_label):
        inputs, outputs, labels = [], [], []
        all_files = os.listdir(path)
        for file in all_files:
            if file == '.DS_Store':
                continue
            file_path = path + file
            with open(file_path, 'r', encoding='utf8') as f:
                data = f.readlines()[1:]
            if labeled is True:
                input_data = [list(map(float, line.strip().split(',')))[
                    :-1] for line in data]
                label = [int(float(line.strip().split(',')[-1]))
                         for line in data]
                if reverse_label:
                    label = [(i - 1) ** 2 for i in label]
            else:
                input_data = [list(map(float, line.strip().split(',')))[
                    :] for line in data]
                label = [0 for _ in data]
            input_np = np.array(input_data)
            time_np = input_np[:, 0] % 86400
            if self.timemode is True:
                input_np[:, 0] = time_np
            mean = input_np.mean(axis=0)  # calculate mean value
            deviation = input_np.std(axis=0)  # calculate standard deviation
            deviation = np.where(deviation != 0, deviation, 1)
            input_np = (input_np - mean) / deviation
            input_data = input_np.tolist()
            input_data, output_data, label = self.cut_data(input_data, label)
            inputs.extend(input_data)
            outputs.extend(output_data)
            labels.extend(label)
        return inputs, labels, outputs

    def cut_data(self, data, label):
        n = 0
        input_data, output_data, anomaly_label = [], [], []
        while n + self.time_feature + self.windowsize <= len(data):
            input_data.append(data[n: n + self.time_feature: self.step])
            output_data.append(
                data[n + self.time_feature: n + self.time_feature + self.windowsize: self.step])
            anomaly_label.append(
                label[n + self.time_feature: n + self.time_feature + self.windowsize: self.step])
            n = n + self.windowsize
        if n + self.time_feature < len(data):
            input_data.append(
                data[-self.time_feature - self.windowsize:-self.windowsize:self.step])
            output_data.append(data[-self.windowsize::self.step])
            anomaly_label.append(label[-self.windowsize::self.step])
        return input_data.copy(), output_data.copy(), anomaly_label.copy()

    def split(self, inputs, labels, outputs):
        train_inputs, train_labels, train_outputs = [], [], []
        test_inputs, test_labels, test_outputs = [], [], []

        split_line = int(len(labels) * (1 - self.test_size))
        train_inputs = inputs[:split_line]
        train_labels = labels[:split_line]
        if len(outputs) > 0:
            train_outputs = outputs[:split_line]
        test_inputs = inputs[split_line:]
        test_labels = labels[split_line:]
        if len(outputs) > 0:
            test_outputs = outputs[split_line:]
        return train_inputs, train_labels, train_outputs, test_inputs, test_labels, test_outputs

    def __len__(self):
        if self.mode == 'train':
            return self.train_inputs.shape[0]
        elif self.mode == 'test':
            return self.test_inputs.shape[0]

    def __getitem__(self, idx):
        if self.mode == 'train':
            return self.train_inputs[idx], self.train_labels[idx], self.train_outputs[idx], idx
        elif self.mode == 'test':
            return self.test_inputs[idx], self.test_labels[idx], self.test_outputs[idx], idx


class GDNAttnDataset(TimeDataset):
    def __init__(self, path, edge_index, input_size=1, windowsize=1, step=1,
                 mode='train', test_size=0, timemode=False, label=False, reverse_label=False):
        super().__init__(path, input_size=input_size, windowsize=windowsize,
                         step=step, mode=mode, test_size=test_size, timemode=timemode,
                         label=label, reverse_label=reverse_label)
        self.edge_index = edge_index

    def __len__(self):
        return super().__len__()

    def __getitem__(self, item):
        if self.mode == 'train':
            return self.train_inputs[item], self.train_labels[item], self.train_outputs[item], item
        elif self.mode == 'test':
            return self.test_inputs[item], self.test_labels[item], self.test_outputs[item], item


if __name__ == "__main__":
    app = TimeDataset('./data/swat/normal/')
    print(app)
